{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9-final"
    },
    "colab": {
      "name": "Robotic Manipulation - Force Control.ipynb",
      "provenance": [],
      "collapsed_sections": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EgiF12Hf1Dhs"
      },
      "source": [
        "**I recommend you run the first code cell of this notebook immediately, to start provisioning drake on the cloud machine, then you can leave this window open as you [read the textbook](manipulation.csail.mit.edu/force.html).**\n",
        "\n",
        "# Notebook setup\n",
        "\n",
        "The following cell will:\n",
        "- on Colab (only), install Drake to `/opt/drake`, install Drake's prerequisites via `apt`, and add pydrake to `sys.path`.  This will take approximately two minutes on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  If you navigate between notebooks using Colab's \"File->Open\" menu, then you can avoid provisioning a separate machine for each notebook.\n",
        "- launch a server for our 3D visualizer (MeshCat) that will be used for the remainder of this notebook.\n",
        "\n",
        "You will need to rerun this cell if you restart the kernel, but it should be fast because the machine will already have drake installed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eeMrMI0-1Dhu"
      },
      "source": [
        "import importlib\n",
        "import sys\n",
        "from urllib.request import urlretrieve\n",
        "\n",
        "# Install drake.\n",
        "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
        "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
        "                \"setup_manipulation_colab.py\")\n",
        "    from setup_manipulation_colab import setup_manipulation\n",
        "    setup_manipulation(manipulation_sha='master', drake_version='20201019', drake_build='nightly')\n",
        "\n",
        "# Determine if this notebook is currently running as a notebook or a unit test.\n",
        "from IPython import get_ipython\n",
        "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
        "\n",
        "# Use pyngrok on colab.\n",
        "server_args = []\n",
        "if 'google.colab' in sys.modules:\n",
        "  server_args = ['--ngrok_http_tunnel']\n",
        "\n",
        "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
        "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
        "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
        "\n",
        "# Let's do all of our imports here, too.\n",
        "import numpy as np\n",
        "import ipywidgets\n",
        "import pydrake.all\n",
        "import os\n",
        "from IPython.display import display\n",
        "\n",
        "from pydrake.all import RigidTransform, RotationMatrix, RollPitchYaw\n",
        "import pydrake.multibody.jupyter_widgets\n",
        "import pydrake.systems.jupyter_widgets\n",
        "from manipulation.utils import FindResource\n",
        "from manipulation.scenarios import SetTransparency, SetColor\n",
        "from manipulation.jupyter_widgets import MakeJointSlidersThatPublishOnCallback"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gOFIUWF2KBPN"
      },
      "source": [
        "def AddPlanarBinAndCracker(plant):\n",
        "    parser = pydrake.multibody.parsing.Parser(plant)\n",
        "    bin = parser.AddModelFromFile(FindResource(\"models/planar_bin.sdf\"))\n",
        "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"bin_base\", bin), RigidTransform(RotationMatrix.MakeZRotation(np.pi/2.0), [0,0,-0.015]))\n",
        "\n",
        "    planar_joint_frame = plant.AddFrame(pydrake.multibody.tree.FixedOffsetFrame(\"planar_joint_frame\", plant.world_frame(), RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
        "    box = parser.AddModelFromFile(pydrake.common.FindResourceOrThrow(\"drake/manipulation/models/ycb/sdf/003_cracker_box.sdf\"))\n",
        "    box_frame = plant.AddFrame(pydrake.multibody.tree.FixedOffsetFrame(\"box_frame\", plant.GetFrameByName(\"base_link_cracker\", box), \n",
        "                                   RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
        "    box_joint = plant.AddJoint(pydrake.multibody.tree.PlanarJoint(\"box_joint\", planar_joint_frame, box_frame))\n",
        "    box_joint.set_default_translation([0, 0.033400])\n",
        "    return box\n",
        "\n",
        "def AddPlanarBinAndSimpleBox(plant, mass=1.0, mu=1.0, width=0.2, depth=0.05, height=0.3):\n",
        "    parser = pydrake.multibody.parsing.Parser(plant)\n",
        "    bin = parser.AddModelFromFile(FindResource(\"models/planar_bin.sdf\"))\n",
        "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"bin_base\", bin), RigidTransform(RotationMatrix.MakeZRotation(np.pi/2.0), [0,0,-0.015]))\n",
        "\n",
        "    planar_joint_frame = plant.AddFrame(pydrake.multibody.tree.FixedOffsetFrame(\"planar_joint_frame\", plant.world_frame(), RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
        "\n",
        "    # TODO(russt): make this a *random* box?\n",
        "    # TODO(russt): move random box to a shared py file.\n",
        "    box_instance = plant.AddModelInstance(\"box\")\n",
        "    box_body = plant.AddRigidBody(\"box_body\", box_instance, pydrake.multibody.tree.SpatialInertia(mass=mass, p_PScm_E=np.array([0., 0., 0.]),\n",
        "            G_SP_E=pydrake.multibody.tree.UnitInertia.SolidBox(width, depth, height)))\n",
        "    if plant.geometry_source_is_registered():\n",
        "        plant.RegisterCollisionGeometry(box_body, RigidTransform(), pydrake.geometry.Box(width-0.001, depth-0.001, height=0.001), \"box_body\", pydrake.multibody.plant.CoulombFriction(mu, mu))\n",
        "        i=0\n",
        "        for x in [-width/2.0, width/2.0]:\n",
        "            for y in [-depth/2.0, depth/2.0]:\n",
        "                for z in [-height/2.0, height/2.0]:\n",
        "                    plant.RegisterCollisionGeometry(box_body, RigidTransform([x, y, z]), pydrake.geometry.Sphere(radius=1e-7), f\"contact_sphere{i}\", pydrake.multibody.plant.CoulombFriction(mu, mu))\n",
        "                    i += 1\n",
        "        plant.RegisterVisualGeometry(box_body, RigidTransform(), pydrake.geometry.Box(width, depth, height), \"box_body\", [.5, .5, .9, 1.0])\n",
        "    box_joint = plant.AddJoint(pydrake.multibody.tree.PlanarJoint(\"box_joint\", planar_joint_frame, box_body.body_frame()))\n",
        "    box_joint.set_default_translation([0, height/2.0])\n",
        "    return box_instance\n",
        "\n",
        "\n",
        "def AddPlanarGripper(plant):\n",
        "    parser = pydrake.multibody.parsing.Parser(plant)\n",
        "    parser.package_map().Add(\"wsg_50_description\", os.path.dirname(pydrake.common.FindResourceOrThrow(\"drake/manipulation/models/wsg_50_description/package.xml\")))\n",
        "    gripper = parser.AddModelFromFile(FindResource(\"models/schunk_wsg_50_welded_fingers.sdf\"), \"gripper\")\n",
        "    gripper_body = plant.GetBodyByName(\"body\", gripper)\n",
        "\n",
        "    # Add a planar joint the old fashioned way (so that I can have three actuators):\n",
        "    gripper_false_body1 = plant.AddRigidBody(\"false_body1\", gripper, pydrake.multibody.tree.SpatialInertia(0, [0,0,0], pydrake.multibody.tree.UnitInertia(0,0,0)))\n",
        "    gripper_false_body2 = plant.AddRigidBody(\"false_body2\", gripper, pydrake.multibody.tree.SpatialInertia(0, [0,0,0], pydrake.multibody.tree.UnitInertia(0,0,0)))\n",
        "    gripper_x = plant.AddJoint(pydrake.multibody.tree.PrismaticJoint(\"gripper_x\", plant.world_frame(), plant.GetFrameByName(\"false_body1\"), [1, 0, 0], -.3, .3))\n",
        "    plant.AddJointActuator(\"gripper_x\", gripper_x)\n",
        "    gripper_z = plant.AddJoint(pydrake.multibody.tree.PrismaticJoint(\"gripper_z\", plant.GetFrameByName(\"false_body1\"), plant.GetFrameByName(\"false_body2\"), [0, 0, 1], 0.0, 0.5))\n",
        "    gripper_z.set_default_translation(0.3)\n",
        "    plant.AddJointActuator(\"gripper_z\", gripper_z)\n",
        "    gripper_frame = plant.AddFrame(pydrake.multibody.tree.FixedOffsetFrame(\"gripper_frame\", plant.GetFrameByName(\"body\", gripper), \n",
        "                                    RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
        "    gripper_theta = plant.AddJoint(pydrake.multibody.tree.RevoluteJoint(\"gripper_theta\", plant.GetFrameByName(\"false_body2\"), gripper_frame, [0, -1, 0],  -np.pi, np.pi))\n",
        "    plant.AddJointActuator(\"gripper_theta\", gripper_theta)\n",
        "\n",
        "    return gripper\n",
        "\n",
        "def MakePlanarGripperOnlyPlant():\n",
        "    plant = pydrake.multibody.plant.MultibodyPlant(time_step=0.005)\n",
        "    AddPlanarGripper(plant)\n",
        "    plant.Finalize()\n",
        "    return plant\n",
        "\n",
        "def AddPointFinger(plant):\n",
        "    mu = 1.0\n",
        "    finger = plant.AddModelInstance(\"finger\")\n",
        "    false_body1 = plant.AddRigidBody(\"false_body1\", finger, pydrake.multibody.tree.SpatialInertia(0, [0,0,0], pydrake.multibody.tree.UnitInertia(0,0,0)))\n",
        "    finger_body = plant.AddRigidBody(\"body\", finger, pydrake.multibody.tree.SpatialInertia(mass=1.0, p_PScm_E=np.array([0., 0., 0.]),\n",
        "            G_SP_E=pydrake.multibody.tree.UnitInertia(1.0, 1.0, 1.0)))\n",
        "    shape = pydrake.geometry.Sphere(0.01)\n",
        "    if plant.geometry_source_is_registered():\n",
        "        plant.RegisterCollisionGeometry(finger_body, RigidTransform(), shape, \"body\", pydrake.multibody.plant.CoulombFriction(mu, mu))\n",
        "        plant.RegisterVisualGeometry(finger_body, RigidTransform(), shape, \"body\", [.9, .5, .5, 1.0])\n",
        "    finger_x = plant.AddJoint(pydrake.multibody.tree.PrismaticJoint(\"finger_x\", plant.world_frame(), plant.GetFrameByName(\"false_body1\"), [1, 0, 0], -.3, .3))\n",
        "    plant.AddJointActuator(\"finger_x\", finger_x)\n",
        "    finger_z = plant.AddJoint(pydrake.multibody.tree.PrismaticJoint(\"finger_z\", plant.GetFrameByName(\"false_body1\"), plant.GetFrameByName(\"body\"), [0, 0, 1], 0.0, 0.5))\n",
        "    finger_z.set_default_translation(0.25)\n",
        "    plant.AddJointActuator(\"finger_z\", finger_z)\n",
        "\n",
        "    return finger\n",
        "\n",
        "def MakeFingerOnlyPlant():\n",
        "    plant = pydrake.multibody.plant.MultibodyPlant(time_step=0.005)\n",
        "    AddPointFinger(plant)\n",
        "    plant.Finalize()\n",
        "    return plant\n",
        "\n",
        "\n",
        "class PointFingerForceControl(pydrake.systems.framework.LeafSystem):\n",
        "\n",
        "    def __init__(self, plant):\n",
        "        pydrake.systems.framework.LeafSystem.__init__(self)\n",
        "        self._plant = plant\n",
        "\n",
        "        self.DeclareVectorInputPort(\n",
        "            \"desired_contact_force\", pydrake.systems.framework.BasicVector(2))\n",
        "        self.DeclareVectorOutputPort(\n",
        "            \"finger_actuation\", pydrake.systems.framework.BasicVector(2), \n",
        "                                    self.CalcOutput)\n",
        "\n",
        "    def CalcOutput(self, context, output):\n",
        "        finger_mass = 1\n",
        "        g = self._plant.gravity_field().gravity_vector()[[0,2]]\n",
        "\n",
        "        desired_force = self.get_input_port(0).Eval(context)\n",
        "        output.SetFromVector(-finger_mass*g - desired_force)\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "source": [
        "# Commanding a constant force"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "def force_control_point_finger(desired_horizontal_forces=[4.0], duration=1.0, draw=False):\n",
        "    builder = pydrake.systems.framework.DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = pydrake.multibody.plant.AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
        "    box = AddPlanarBinAndCracker(plant)\n",
        "    finger = AddPointFinger(plant)\n",
        "    plant.GetJointByName(\"finger_x\").set_default_translation(0.15)\n",
        "    plant.GetJointByName(\"finger_z\").set_default_translation(0.05)\n",
        "    plant.Finalize()\n",
        "\n",
        "    if draw:\n",
        "        vis = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(builder, scene_graph, delete_prefix_on_load=False)\n",
        "        vis.vis['drake']['4'].delete()  # Only flush the gripper geometry.\n",
        "    #    vis.set_planar_viewpoint(xmin=-.5, xmax=.5, ymin=-0.1, ymax=0.6)\n",
        "        contact_visualizer = builder.AddSystem(pydrake.systems.meshcat_visualizer.MeshcatContactVisualizer(vis, plant=plant, contact_force_radius=0.005))\n",
        "        builder.Connect(plant.get_contact_results_output_port(),\n",
        "                        contact_visualizer.GetInputPort(\"contact_results\"))\n",
        "\n",
        "    controller = builder.AddSystem(PointFingerForceControl(plant))\n",
        "    builder.Connect(controller.get_output_port(), plant.get_actuation_input_port())\n",
        "    builder.ExportInput(controller.get_input_port(), \"desired_contact_force\")\n",
        "\n",
        "    log = pydrake.systems.primitives.LogOutput(plant.get_state_output_port(finger), builder)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "\n",
        "    plt.figure()\n",
        "    lines = []\n",
        "    for f in desired_horizontal_forces:\n",
        "        log.reset()\n",
        "        simulator = pydrake.systems.analysis.Simulator(diagram)\n",
        "        context = simulator.get_mutable_context()\n",
        "        diagram.get_input_port().FixValue(context, [f, 0])\n",
        "\n",
        "        simulator.AdvanceTo(duration)\n",
        "        t = log.sample_times()\n",
        "        x = log.data()\n",
        "        lines.append(plt.plot(t, x[0, :])[0])\n",
        "\n",
        "    plt.legend(lines, [(\"$f^{c_d}$ = \" + str(f)) for f in desired_horizontal_forces])\n",
        "    plt.xlabel('t (sec)')\n",
        "    plt.ylabel('x (m)')\n",
        "\n",
        "force_control_point_finger([0.25, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VBc7zDKlElKu"
      },
      "source": [
        "#  A force-based flip-up strategy"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "tags": [],
        "id": "xJ7sDtzbEfYw"
      },
      "source": [
        "from ipywidgets import FloatSlider, Layout\n",
        "\n",
        "def force_control_point_finger(teleop=True):\n",
        "    builder = pydrake.systems.framework.DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = pydrake.multibody.plant.AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
        "    box = AddPlanarBinAndCracker(plant)\n",
        "    finger = AddPointFinger(plant)\n",
        "    plant.GetJointByName(\"finger_x\").set_default_translation(0.15)\n",
        "    plant.GetJointByName(\"finger_z\").set_default_translation(0.025)\n",
        "    plant.Finalize()\n",
        "\n",
        "    vis = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(builder, scene_graph, delete_prefix_on_load=False)\n",
        "    vis.vis['drake']['4'].delete()  # Only flush the gripper geometry.\n",
        "    #vis.set_planar_viewpoint(xmin=-.5, xmax=.5, ymin=-0.1, ymax=0.6)\n",
        "    contact_visualizer = builder.AddSystem(pydrake.systems.meshcat_visualizer.MeshcatContactVisualizer(vis, plant=plant, contact_force_radius=0.005, contact_force_scale=40.0))\n",
        "    builder.Connect(plant.get_contact_results_output_port(),\n",
        "                    contact_visualizer.GetInputPort(\"contact_results\"))\n",
        "\n",
        "    class RotateBoxOnTopCorner(pydrake.systems.framework.LeafSystem):\n",
        "\n",
        "        def __init__(self, plant):\n",
        "            pydrake.systems.framework.LeafSystem.__init__(self)\n",
        "            self._plant = plant\n",
        "\n",
        "            # One state variable for the integral error.\n",
        "            self.DeclareContinuousState(1)\n",
        "\n",
        "            self.DeclareVectorInputPort(\"desired_box_angle\",\n",
        "                pydrake.systems.framework.BasicVector(1))\n",
        "            self.DeclareVectorInputPort(\n",
        "                \"box_state\", pydrake.systems.framework.BasicVector(6))\n",
        "            self.DeclareVectorInputPort(\n",
        "                \"finger_state\", pydrake.systems.framework.BasicVector(4))\n",
        "            self.DeclareVectorOutputPort(\n",
        "                \"finger_force\", pydrake.systems.framework.BasicVector(2), \n",
        "                                        self.CalcOutput)\n",
        "\n",
        "        def DoCalcTimeDerivatives(self, context, derivatives):\n",
        "            # Compute integral of the angle error.\n",
        "            theta_desired = self.get_input_port(0).Eval(context)[0]\n",
        "            theta = self.get_input_port(1).Eval(context)[2]\n",
        "            derivatives.get_mutable_vector().SetFromVector([theta_desired - theta])\n",
        "\n",
        "        def CalcOutput(self, context, output):\n",
        "            box_mass = 0.411  # Let's revisit whether we need this!\n",
        "            mu_A = 1.0\n",
        "            mu_C = 1.0\n",
        "            g = -self._plant.gravity_field().gravity_vector()[2]\n",
        "\n",
        "            theta_desired = self.get_input_port(0).Eval(context)[0]\n",
        "            box_state = self.get_input_port(1).Eval(context)\n",
        "            theta = box_state[2]\n",
        "            theta_dot = box_state[5]\n",
        "\n",
        "            R_WC = np.array([[-np.sin(theta), -np.cos(theta)], [np.cos(theta), -np.sin(theta)]])\n",
        "            R_CW = R_WC.T\n",
        "            finger_state = self.get_input_port(2).Eval(context)\n",
        "\n",
        "            prog = pydrake.solvers.mathematicalprogram.MathematicalProgram()\n",
        "            f_C_W = prog.NewContinuousVariables(2, \"f_C_W\")  # force from finger to box\n",
        "            f_C_C = np.matmul(R_CW, f_C_W)\n",
        "\n",
        "            # PID control on f_CC[1]\n",
        "            kp = 100.0\n",
        "            ki = 1.0\n",
        "            kd = 2.0*np.sqrt(kp)  # should really be kp*I on inside\n",
        "            integral_error = context.get_continuous_state_vector().GetAtIndex(0)\n",
        "            pid = kp*(theta_desired - theta) + ki*integral_error - kd*theta_dot\n",
        "            prog.AddCost((f_C_C[0] - pid)**2)\n",
        "\n",
        "            prog.AddConstraint(f_C_W[0] <= mu_A * (box_mass*g - f_C_W[1]))\n",
        "            prog.AddConstraint(-f_C_W[0] <= mu_A * (box_mass*g - f_C_W[1]))\n",
        "            prog.AddConstraint(f_C_C[1] >= 0.1)  # To ensure I make and stay in contact.\n",
        "            prog.AddConstraint(f_C_C[0] <= mu_C * f_C_C[1])\n",
        "            prog.AddConstraint(-f_C_C[0] <= mu_C * f_C_C[1])\n",
        "\n",
        "            result = pydrake.solvers.mathematicalprogram.Solve(prog)\n",
        "            \n",
        "            if not result.is_success():\n",
        "                print(result.GetInfeasibleConstraintNames(prog))\n",
        "                output.SetFromVector([0, 0])\n",
        "                return\n",
        "            \n",
        "            f_C_W = result.GetSolution(f_C_W)\n",
        "            output.SetFromVector(-f_C_W)\n",
        "\n",
        "\n",
        "    controller = builder.AddSystem(PointFingerForceControl(plant))\n",
        "    builder.Connect(controller.get_output_port(), plant.get_actuation_input_port())\n",
        "\n",
        "    rotate_box = builder.AddSystem(RotateBoxOnTopCorner(plant))\n",
        "    builder.Connect(rotate_box.get_output_port(), controller.get_input_port())\n",
        "    builder.Connect(plant.get_state_output_port(box), rotate_box.get_input_port(1))\n",
        "    builder.Connect(plant.get_state_output_port(finger), rotate_box.get_input_port(2))\n",
        "\n",
        "    if teleop:\n",
        "        slider = FloatSlider(value=0, min=0, max=np.pi/2, step=0.05, continuous_update=True, description=\"desired_box_angle\", layout=Layout(width=\"'100'\"))\n",
        "        teleop = builder.AddSystem(pydrake.systems.jupyter_widgets.WidgetSystem([slider]))\n",
        "        builder.Connect(teleop.get_output_port(), rotate_box.get_input_port(0)) \n",
        "    else:\n",
        "        theta_trajectory = pydrake.trajectories.PiecewisePolynomial.FirstOrderHold(\n",
        "            [0, 2, 20], [[0, 0, np.pi/2.0]])\n",
        "        trajectory_source = builder.AddSystem(pydrake.systems.primitives.TrajectorySource(theta_trajectory))\n",
        "        builder.Connect(trajectory_source.get_output_port(), rotate_box.get_input_port(0))\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    simulator = pydrake.systems.analysis.Simulator(diagram)\n",
        "\n",
        "    if not running_as_notebook:\n",
        "        simulator.AdvanceTo(0.01)\n",
        "    elif teleop:\n",
        "        stop_button = ipywidgets.ToggleButton(value=False, description='Stop Simulation')\n",
        "        display(stop_button)\n",
        "        while not stop_button.value:\n",
        "            simulator.AdvanceTo(simulator.get_context().get_time() + 2.0)\n",
        "        stop_button.value = False\n",
        "    else:\n",
        "        simulator.set_target_realtime_rate(.5)  # slow motion!\n",
        "        vis.start_recording()\n",
        "        simulator.AdvanceTo(theta_trajectory.end_time())\n",
        "        vis.stop_recording()\n",
        "        vis.vis['drake/contact_forces'].delete()  # TODO: make my recording include the contact forces!\n",
        "        vis.publish_recording()\n",
        "\n",
        "force_control_point_finger(teleop=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7q0A14bAilIX"
      },
      "source": [
        "# Teleop with Stiffness Control\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JX1Oq7gfWsBD"
      },
      "source": [
        "def teleop_box_pickup(point_finger=True):\n",
        "    builder = pydrake.systems.framework.DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = pydrake.multibody.plant.AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
        "    AddPlanarBinAndSimpleBox(plant)\n",
        "#    AddPlanarBinAndCracker(plant)\n",
        "    SetTransparency(scene_graph, alpha=0.5, source_id=plant.get_source_id())\n",
        "\n",
        "    gripper = AddPointFinger(plant) if point_finger else AddPlanarGripper(plant)\n",
        "    plant.Finalize()\n",
        "\n",
        "    vis = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(builder, scene_graph)\n",
        "    vis.set_planar_viewpoint(xmin=-.5, xmax=.5, ymin=-0.1, ymax=0.6)\n",
        "    contact_visualizer = builder.AddSystem(pydrake.systems.meshcat_visualizer.MeshcatContactVisualizer(vis, plant=plant, contact_force_radius=0.005, contact_force_scale=40.0))\n",
        "    builder.Connect(plant.get_contact_results_output_port(),\n",
        "                    contact_visualizer.GetInputPort(\"contact_results\"))\n",
        "\n",
        "\n",
        "    # Use the controller plant to visualize the set point geometry.\n",
        "    controller_scene_graph = builder.AddSystem(pydrake.geometry.SceneGraph())\n",
        "    controller_plant = pydrake.multibody.plant.MultibodyPlant(time_step=0.005)\n",
        "    controller_plant.RegisterAsSourceForSceneGraph(controller_scene_graph)\n",
        "    AddPointFinger(controller_plant) if point_finger else AddPlanarGripper(controller_plant)\n",
        "    controller_plant.Finalize()\n",
        "    SetColor(controller_scene_graph, color=[1.0, 165.0/255, 0.0, 1.0], source_id=plant.get_source_id())\n",
        "    controller_vis = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(builder, controller_scene_graph, prefix=\"controller\", delete_prefix_on_load=False)\n",
        "    controller_vis.set_name(\"controller meshcat\")\n",
        "    controller_vis.vis['controller']['2'].delete()  # Only flush the gripper geometry.\n",
        "\n",
        "    # TODO: Implement the proper JointStiffnessController in drake and use it here.\n",
        "    # (For a point finger with unit mass, the InverseDynamicsController is identical)\n",
        "    N = controller_plant.num_positions()\n",
        "    kp = [100]*N\n",
        "    ki = [1]*N\n",
        "    kd = [2*np.sqrt(kp[0])]*N\n",
        "    controller = builder.AddSystem(pydrake.systems.controllers.InverseDynamicsController(controller_plant, kp, ki, kd, False))\n",
        "    builder.Connect(plant.get_state_output_port(gripper), controller.get_input_port_estimated_state())\n",
        "\n",
        "    display(ipywidgets.Label(value=\"Set desired positions:\"))\n",
        "    position_sliders = builder.AddSystem(pydrake.multibody.jupyter_widgets.JointSliders(controller_plant))\n",
        "    positions_to_state = builder.AddSystem(pydrake.systems.primitives.Multiplexer([N, N]))\n",
        "    builder.Connect(position_sliders.get_output_port(), positions_to_state.get_input_port(0))\n",
        "    zeros = builder.AddSystem(pydrake.systems.primitives.ConstantVectorSource([0]*N))\n",
        "    builder.Connect(zeros.get_output_port(0), positions_to_state.get_input_port(1))\n",
        "    builder.Connect(positions_to_state.get_output_port(), controller.get_input_port_desired_state())\n",
        "\n",
        "    positions_to_poses = builder.AddSystem(pydrake.systems.rendering.MultibodyPositionToGeometryPose(controller_plant))\n",
        "    builder.Connect(position_sliders.get_output_port(), positions_to_poses.get_input_port())\n",
        "    builder.Connect(positions_to_poses.get_output_port(), controller_scene_graph.get_source_pose_port(controller_plant.get_source_id()))\n",
        "    builder.Connect(controller.get_output_port(), plant.get_actuation_input_port())\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    simulator = pydrake.systems.analysis.Simulator(diagram)\n",
        "\n",
        "    if running_as_notebook:\n",
        "        stop_button = ipywidgets.ToggleButton(value=False, description='Stop Simulation')\n",
        "        display(stop_button)\n",
        "        while not stop_button.value:\n",
        "            simulator.AdvanceTo(simulator.get_context().get_time() + 2.0)\n",
        "        stop_button.value = False\n",
        "    else:\n",
        "        simulator.AdvanceTo(0.01)\n",
        "\n",
        "teleop_box_pickup(point_finger=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "source": [
        "# A stiffness-control-based flip-up strategy"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xLM8906QEfY0"
      },
      "source": [
        "def stiffness_box_flipup():\n",
        "    builder = pydrake.systems.framework.DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = pydrake.multibody.plant.AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
        "    AddPlanarBinAndSimpleBox(plant)\n",
        "    SetTransparency(scene_graph, alpha=0.5, source_id=plant.get_source_id())\n",
        "\n",
        "    gripper = AddPointFinger(plant)\n",
        "    plant.GetJointByName(\"finger_x\").set_default_translation(0.15)\n",
        "    plant.GetJointByName(\"finger_z\").set_default_translation(0.025)\n",
        "    plant.Finalize()\n",
        "\n",
        "    vis = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(builder, scene_graph)\n",
        "    vis.set_planar_viewpoint(xmin=-.5, xmax=.5, ymin=-0.1, ymax=0.6)\n",
        "    contact_visualizer = builder.AddSystem(pydrake.systems.meshcat_visualizer.MeshcatContactVisualizer(vis, plant=plant, contact_force_radius=0.005, contact_force_scale=40.0))\n",
        "    builder.Connect(plant.get_contact_results_output_port(),\n",
        "                    contact_visualizer.GetInputPort(\"contact_results\"))\n",
        "\n",
        "    # Use the controller plant to visualize the set point geometry.\n",
        "    controller_scene_graph = builder.AddSystem(pydrake.geometry.SceneGraph())\n",
        "    controller_plant = pydrake.multibody.plant.MultibodyPlant(time_step=0.005)\n",
        "    controller_plant.RegisterAsSourceForSceneGraph(controller_scene_graph)\n",
        "    AddPointFinger(controller_plant)\n",
        "    controller_plant.Finalize()\n",
        "    SetColor(controller_scene_graph, color=[1.0, 165.0/255, 0.0, 1.0], source_id=plant.get_source_id())\n",
        "    controller_vis = pydrake.systems.meshcat_visualizer.ConnectMeshcatVisualizer(builder, controller_scene_graph, prefix=\"controller\", delete_prefix_on_load=False)\n",
        "    controller_vis.set_name(\"controller meshcat\")\n",
        "    controller_vis.vis['controller']['2'].delete()  # Only flush the gripper geometry.\n",
        "\n",
        "    # TODO: Implement the proper JointStiffnessController in drake and use it here.\n",
        "    # (For a point finger with unit mass, the InverseDynamicsController is identical)\n",
        "    N = controller_plant.num_positions()\n",
        "    kp = [100]*N\n",
        "    ki = [1]*N\n",
        "    kd = [2*np.sqrt(kp[0])]*N\n",
        "    controller = builder.AddSystem(pydrake.systems.controllers.InverseDynamicsController(controller_plant, kp, ki, kd, False))\n",
        "    builder.Connect(plant.get_state_output_port(gripper), controller.get_input_port_estimated_state())\n",
        "\n",
        "    traj = pydrake.trajectories.PiecewisePolynomial.FirstOrderHold([0, 10.0, 20.0], np.array([[.15, .025], [-.25, .025], [-.25, 0.125]]).T)\n",
        "    traj_system = builder.AddSystem(pydrake.systems.primitives.TrajectorySource(traj))\n",
        "    positions_to_state = builder.AddSystem(pydrake.systems.primitives.Multiplexer([N, N]))\n",
        "    builder.Connect(traj_system.get_output_port(), positions_to_state.get_input_port(0))\n",
        "    zeros = builder.AddSystem(pydrake.systems.primitives.ConstantVectorSource([0]*N))\n",
        "    builder.Connect(zeros.get_output_port(0), positions_to_state.get_input_port(1))\n",
        "    builder.Connect(positions_to_state.get_output_port(), controller.get_input_port_desired_state())\n",
        "\n",
        "    positions_to_poses = builder.AddSystem(pydrake.systems.rendering.MultibodyPositionToGeometryPose(controller_plant))\n",
        "    builder.Connect(traj_system.get_output_port(), positions_to_poses.get_input_port())\n",
        "    builder.Connect(positions_to_poses.get_output_port(), controller_scene_graph.get_source_pose_port(controller_plant.get_source_id()))\n",
        "    builder.Connect(controller.get_output_port(), plant.get_actuation_input_port())\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    simulator = pydrake.systems.analysis.Simulator(diagram)\n",
        "\n",
        "    if running_as_notebook:\n",
        "        simulator.AdvanceTo(traj.end_time()+0.5)\n",
        "    else:\n",
        "        simulator.AdvanceTo(0.01)\n",
        "\n",
        "stiffness_box_flipup()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g3lNMaloEfY3"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}